import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;

/**
 * 378. 有序矩阵中第K小的元素
 * https://leetcode-cn.com/problems/kth-smallest-element-in-a-sorted-matrix/
 */
public class Solutions_378 {
    public static void main(String[] args) {
        int[][] matrix = {{1, 5, 9},
                          {10, 11, 13},
                          {12, 13, 15}};
        int k = 8;  // output: 13

        int result = kthSmallest(matrix, k);
        System.out.println(result);
    }

    /**
     * 解法三：二分法（ms）
     */
    public static int kthSmallest(int[][] matrix, int k) {
        int n = matrix.length;
        int left = matrix[0][0];
        int right = matrix[n - 1][n - 1];
        while (left < right) {
            int mid = left + (right - left) / 2;
            // matrix 中，小于等于 mid 的个数
            int count = check(matrix, mid, k, n);

            if (count >= k) {
                // 舍弃右下角区域
                right = mid;
            } else {
                // 舍弃左上角区域
                left = mid + 1;
            }
        }
        return left;
    }

    public static int check(int[][] matrix, int mid, int k, int n) {
        int i = n - 1;
        int j = 0;
        int count = 0;
        while (i >= 0 && j < n) {
            if (matrix[i][j] <= mid) {
                // 当前列中，小于等于 mid 的元素个数
                count += i + 1;
                // 列扩大范围，数值更大
                j ++;
            } else {
                // 行缩小范围，数值更小
                i --;
            }
        }
        // matrix 中，小于等于 mid 的元素个数
        return count;
    }

    /**
     * 解法二：归并排序（23ms）
     */
    public static int kthSmallest3(int[][] matrix, int k) {
        PriorityQueue<int[]> queue = new PriorityQueue<int[]>(new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                // 排序规则：升序
                return o1[0] - o2[0];
            }
        });
        int row = matrix.length;
        // 将每行第一个元素入队，索引 0：数值，索引 1：行，索引 2：列
        for (int i = 0; i < row; i++) {
            queue.offer(new int[]{matrix[i][0], i, 0});
        }
        // 遍历 k - 1 次，弹出前 k - 1 个较小的元素，最终 queue 中首位即是排序后第 k 小元素
        for (int i = 0; i < k - 1; i++) {
            // 最小的元素
            int[] arr = queue.poll();
            // 最小元素所在行的下一元素入队
            int m = arr[1], n = arr[2];
            if (n != row - 1) {
                queue.offer(new int[]{matrix[m][n + 1], m, n + 1});
            }
        }
        return queue.poll()[0];
    }

    /**
     * 解法一：排序（8ms）
     */
    public static int kthSmallest2(int[][] matrix, int k) {
        int row = matrix.length, col = matrix[0].length;
        // 二维数组转一维数组
        int[] res = new int[row * col];
        for (int i = 0; i < row; i++) {
            for (int j = 0; j < col; j++) {
                res[i * row + j] = matrix[i][j];
            }
        }
        // 排序后，取第 k 个
        Arrays.sort(res);
        return res[k - 1];
    }
}
